#  -*- coding: utf-8 -*-

import sys
sys.path.append('/Users/xbs/Code/HunterQuant')
from pymongo import UpdateOne
from util.database import DB_CONN
from data.data_module import DataModule
from util.stock_util import get_trading_dates,multi_computer
from datetime import datetime
import time
import pandas as pd
import os
from util.database import base_data_path

"""
从tushare获取日K数据，保存到本地的MongoDB数据库中
"""
index_nofq_path = f"{base_data_path}/index_nofq"
code_nofq_path = f"{base_data_path}/code_nofq"
code_qfq_path = f"{base_data_path}/code_qfq"
code_hfq_path = f"{base_data_path}/code_hfq"

class DailyCrawler:
    def __init__(self):
        self.daily = DB_CONN['daily']
        self.daily_qfq = DB_CONN['daily_qfq']
        self.daily_hfq = DB_CONN['daily_hfq']
        self.dm = DataModule()


    def crawl_index(self, begin_date=None, end_date=None):
        """
        抓取指数的日线数据，并保存到本地数据数据库中
        抓取的日期范围从A股成立至今

        指数代码：上证指数：999999 深圳成指：399001，中小板指：399005 创业板指：399006
                  深圳综指：399106 深圳100R：399004 深圳300：399007 中小300：399008
                  上证50：000016 上证180：000010 上证380：000009 沪深300：000300
                  中证100：000903 中证500：000905
        """
        index_codes = self.tdx_get_all_indexes()

        # 设置默认的日期范围,A股成立时间
        if begin_date is None:
            begin_date = '1990-12-19'

        if end_date is None:
            end_date = datetime.now().strftime('%Y-%m-%d')

        index_num = len(index_codes)
        i = 0
        #建立code+date的索引，提高save_data时写入数据的查询速度
        self.daily.create_index([('code', 1), ('date', 1),('index',1)])

        #优先处理上证指数，获得交易日期
        df_daily = self.tdx_get_k_data("999999", index=True, start=begin_date, end=end_date)
        self.save_data("999999", df_daily, self.daily, {'index': True})

        for code in index_codes:
  
            start_time = time.time()

            df_daily = self.tdx_get_k_data(code, index=True, start=begin_date, end=end_date)
            self.save_data(code, df_daily, self.daily, {'index': True})

            end_time = time.time()
            i = i + 1
            print("已完成处理的指数个数：%4d, 所有的指数个数：%4d, 本次执行时间: %.3f 秒 " % (i, index_num, (end_time-start_time)), flush=True)

    def save_data(self, code, df_daily, collection, extra_fields=None):
        """
        将从网上抓取的数据保存到本地MongoDB中

        :param code: 股票代码
        :param df_daily: 包含日线数据的DataFrame
        :param collection: 要保存的数据集
        :param extra_fields: 除了K线数据中保存的字段，需要额外保存的字段
        """
        update_requests = []
        close_tmp = 0
        for df_index in df_daily.index:
            daily_obj = df_daily.loc[df_index]

            doc = self.daily_obj_2_doc(code, daily_obj)
            # 日期是df的index,所以单独赋值
            doc["date"] = df_index

            if 0 == doc['volume'] :
                print("股票：%s,日期：%s,成交量为0，为停牌日，当日数据不处理，后续修复" % (code,df_index))
                continue

            #指数情况，直接填充pre_close,首日的pre_close用开盘价填充
            is_index = False
            if extra_fields == {'index': True}:
                is_index = True
                if close_tmp == 0:
                    close_tmp = doc['open']
                doc.update({'pre_close': close_tmp})
                close_tmp = doc['close']
            doc.update(extra_fields)

            is_insert = True
            if is_index:
                df_index_daily = self.dm.get_k_data(code, index=is_index, autype=None, begin_date=doc['date'], end_date=doc['date'])
                #数据库里找不到数据才插入
                if df_index_daily.index.size > 0:
                    is_insert = False

            if is_insert:
                update_requests.append(
                    UpdateOne(
                        {'code': doc['code'], 'date': doc['date'], 'index': doc['index']},
                        {'$set': doc},
                        upsert=True)
                )
        # 批量写入，提高访问效率
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存日线数据到数据集：%s，代码： %s, 插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, code, update_result.upserted_count, update_result.modified_count,(end_time-start_time)),
                  flush=True)

    def multi_crawl_codes(self,begin_date=None, end_date=None):
        codes = self.tdx_get_all_codes()
        args = (begin_date,end_date)
        multi_computer(crawl,codes,args)

    @staticmethod
    def daily_obj_2_doc(code, daily_obj):
        return {
            'code': code,
            'close': daily_obj['    收盘'],
            'open': daily_obj['    开盘'],
            'high': daily_obj['    最高'],
            'low': daily_obj['    最低'],
            'volume': daily_obj['    成交量'],
            'amount': daily_obj['    成交额'],
            'is_trading':True
        }

    @staticmethod
    def tdx_get_k_data(code=None, start=None, end=None,
                  ktype='D', autype=None,
                  index=False):
        """
            获取k线数据
            ---------
            Parameters:
              code:string
                          股票代码 e.g. 600848
              start:string
                          开始日期 format：YYYY-MM-DD 为空时取上市首日
              end:string
                          结束日期 format：YYYY-MM-DD 为空时取最近一个交易日
              autype:string
                          复权类型，qfq-前复权 hfq-后复权 None-不复权，默认为qfq
              ktype：string
                          数据类型，D=日k线 W=周 M=月 5=5分钟 15=15分钟 30=30分钟 60=60分钟，默认为D
            return
            -------
              DataFrame
                  date 交易日期(index)
                  open 开盘价
                  high  最高价
                  close 收盘价
                  low 最低价
                  volume 成交量
                  amount 成交额
            """
        filename = code + ".txt"
        print(f"获取k线原始数据文件：代码：{code},复权类型：{autype},是否指数：{index}")
        empty_df = pd.DataFrame()
        #目前只处理日线数据
        if ktype != "D":
            pass

        if index:
            file_path = os.path.join(index_nofq_path,filename)
        else:
            if autype == "qfq":
                file_path = os.path.join(code_qfq_path,filename)
            elif autype == "hfq":
                file_path = os.path.join(code_hfq_path,filename)
            elif autype == None:
                file_path = os.path.join(code_nofq_path,filename)
            else:
                file_path = None
        try:
            data_df = pd.read_table(file_path, encoding = "gb2312", skiprows=[0])
            result_df = data_df.set_index('      日期')
            result_df.drop("数据来源:通达信",inplace=True) #删除最后一行，并且替换原数据
        except:
            print("打开历史数据文件错误，错误文件是：%s" % file_path, flush=True)
            return empty_df
        #通过获取交易日期，处理不存在的交易日入参
        try:
            #无起止时间或者是上证指数时，返回全部数据，上证指数数据用于记录交易日期
            if (start is None) & (end is None) or (index == True and code == "999999"):
                return result_df
            trade_dates = get_trading_dates(start,end)
            trade_dates_set = set(trade_dates)
            index_set = set(result_df.index)
            #如果传入的交易日期与文件内的交易日期无交集，则返回全部交易数据;如果存在交集，则返回交集数据
            if True == trade_dates_set.isdisjoint(index_set):
                ret_set = index_set
            else:
                ret_set = trade_dates_set & index_set
            #set是乱序的，重组为lst，并重新排序
            ret_lst = list(ret_set)
            ret_lst.sort()
            ret_df = result_df.loc[ret_lst,:]

            return ret_df

        except:
            print("获取指定日期范围的DataFrame错误，错误文件是：%s, 错误日期范围是：%s : %s" % (file_path, start, end), flush=True)
            return empty_df

    @staticmethod
    def tdx_get_all_codes():
        all_codes = [x.strip(".txt") for x in os.listdir(code_nofq_path)]
        return all_codes

    @staticmethod
    def tdx_get_all_indexes():
        all_codes = [x.strip(".txt") for x in os.listdir(index_nofq_path)]
        return all_codes

def crawl(codes=None, begin_date=None, end_date=None):
    """
    获取所有股票从A股成立以来至今的K线数据（包括前复权、后复权和不复权三种），保存到数据库中
    """
    self = DailyCrawler()
    if codes is None:
        # 获取所有股票代码
        codes = self.tdx_get_all_codes()

    codes_num = len(codes)


    # 设置默认的日期范围,tdx有数据的时间
    if begin_date is None:
        begin_date = '1995-02-24'

    if end_date is None:
        end_date = datetime.now().strftime('%Y-%m-%d')

    i = 0
    #建立code+date的索引，提高save_data时写入数据的查询速度

    self.daily.create_index([('code', 1), ('date', 1)])
    self.daily_qfq.create_index([('code', 1), ('date', 1)])
    self.daily_hfq.create_index([('code', 1), ('date', 1)])

    for code in codes:
        # 抓取不复权的价格
        start_time = time.time()
        df_daily = self.tdx_get_k_data(code, autype=None, start=begin_date, end=end_date)
        if not df_daily.empty:
            self.save_data(code, df_daily, self.daily, {'index': False})

        # 抓取前复权的价格,前复权每天都全部更新
        df_daily_qfq = self.tdx_get_k_data(code, autype='qfq')
        if not df_daily_qfq.empty:
            self.save_data(code, df_daily_qfq, self.daily_qfq, {'index': False})

        # 抓取后复权的价格
        df_daily_hfq = self.tdx_get_k_data(code, autype='hfq', start=begin_date, end=end_date)
        if not df_daily_hfq.empty:
            self.save_data(code, df_daily_hfq, self.daily_hfq, {'index': False})
        end_time = time.time()
        i = i +1
        print("已完成处理的股票个数：%4d, 所有的股票个数：%4d, 本次执行时间: %.3f 秒 " % (i, codes_num, (end_time-start_time)), flush=True)

if __name__ == '__main__':
    dc = DailyCrawler()
    dc.crawl_index(begin_date='2020-11-12',end_date='2020-11-13')
    #crawl()
    #retdf = dc.tdx_get_k_data("399001",start='1995-02-02', end='1995-02-28', index=True)
    #DB_CONN['daily'].create_index([('code',1),('date',1)])
    #doc = {"code":"000001","date":"2018-01-08",'index': False,"close":"12.96",'open':"13.25","high":13.29,"low":12.86,"volume":21586200, "amount":2806098944.00}
    #t0 = time.time()
    #print(DB_CONN['daily'].update_one({'code': "000001", 'date': "2018-01-08", 'index': False },{'$set': doc},upsert=True),time.time() - t0)

